package com.zzw.websocket;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.net.URI;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author Administrator
 */
public class DefaultTextHandler extends TextWebSocketHandler {

    private static final Logger logger = LoggerFactory.getLogger(DefaultTextHandler.class);

    private static final ConcurrentHashMap<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, String> TOKEN_SESSION = new ConcurrentHashMap<>();

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        logger.error("session transport error", exception);
        super.handleTransportError(session, exception);
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        URI uri = session.getUri();
        assert uri != null;
        String[] paths = uri.getPath().split("/");
        String token = paths[paths.length - 1];
        if (!Objects.equals(token, "123456")) {
            session.close(CloseStatus.NOT_ACCEPTABLE);
            return;
        }
        sessions.put(token, session);
        TOKEN_SESSION.put(session.getId(), token);
        session.sendMessage(new TextMessage(token + ",welcome!"));
        super.afterConnectionEstablished(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        if (!TOKEN_SESSION.containsKey(session.getId())) {
            return;
        }
        if (TOKEN_SESSION.containsKey(TOKEN_SESSION.get(session.getId()))) {
            sessions.remove(TOKEN_SESSION.get(session.getId()));
        }
        super.afterConnectionClosed(session, status);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        logger.info("session:{},token:{},message:{}", session.getId(), TOKEN_SESSION.get(session.getId()), message.getPayload());
        super.handleTextMessage(session, message);
    }
}
